import json
import time
import re
import os
from openai import OpenAI
import fitz  # PyMuPDF
from docx import Document

# ========================== 1. 配置区 ==========================
API_KEY = "BearerCubJNNWPCQjooxiEbLGB:IgSlCSEPDRuFXHJlXlMw" 
URL = "https://spark-api-open.xf-yun.com/v2/chat/completions"
INPUT_DOCUMENT_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/豫洛阳-兰郑长干线-CPY-0790-BFGDGS-ZZSYQFGS.docx"
OUTPUT_JSON_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/image_text_consistency.json"
SAMPLES_PER_TYPE = 3  # 每种类型生成的样本数量（4种类型 × 3个 = 12个负样本）


# ========================== 2. 功能函数区 ==========================
def get_answer(message):
    """改为使用 OpenAI 兼容接口（binapi.shop），非流式聚合输出"""
    try:
        client = OpenAI(
            api_key="sk-ma0HIkN3beK0mtThtDzIfRLN5IzpL4agmB9RSbpHk53h7BF5",
            base_url="https://www.chataiapi.com/v1",
        )
        response = client.chat.completions.create(
            model="gemini-2.5-pro",
            messages=message,
            temperature=0.7,
        )
        content = response.choices[0].message.content or ""
        if content.strip():
            print("\n\n*******************模型生成内容如下********************\n")
            print(content)
        return content
    except Exception as e:
        print(f"API请求失败：{str(e)}")
        return ""


def read_document(file_path):
    """
    读取 Word 文件，仅提取表格中"高后果区特征描述"行的内容。
    """
    try:
        doc = Document(file_path)
        feature_text = ""
        # ========= ① 提取表格中的“高后果区特征描述”行 =========
        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    if "高后果区特征描述" in cell.text:
                        # 把整行内容取出来
                        feature_text = ' | '.join([c.text.strip() for c in row.cells if c.text.strip()])
                        break
                if feature_text:
                    break
            if feature_text:
                break

        if not feature_text:
            return "错误: 未找到【高后果区特征描述】行"

        # 直接返回表格中的特征描述内容
        return feature_text

    except Exception as e:
        return f"错误: 读取文件失败 {file_path} - {str(e)}"



# ========================== 负样本模板（保持不变） ==========================
def get_negative_prompt(text_chunk, source_label, prompt_type):
    prompt_templates = {
        1: f"""任务：为【{source_label}】生成“文本新增建筑（图中无）”类型的负样本。
说明：
- 假设原始文本中图片<image>的标注是正确的。
- 请基于原文中真实出现的地名、建筑或小区名进行改写。
- 在文本中**新增**一个在图<image>中未标注的建筑或地点（例如增加“中房朗悦湾”“某村”等在文档同区段内出现过的实体）。
原始文本：
---
{text_chunk}
---
输出：仅返回修改后的文本。""",

        2: f"""任务：为【{source_label}】生成“文本漏写建筑（图中有）”类型的负样本。
说明：
- 假设图片<image>标注了多个建筑物或地点。
- 请删除部分在原文中出现的建筑描述（如删掉一个或两个小区名）。
原始文本：
---
{text_chunk}
---
输出：仅返回修改后的文本。""",

        3: f"""任务：为【{source_label}】生成“建筑物名称不一致”类型的负样本。
说明：
- 修改文本中已有的建筑名称，使其与原文出现的名称存在轻微差异。
原始文本：
---
{text_chunk}
---
输出：仅返回修改后的文本。""",

        4: f"""任务：为【{source_label}】生成“方位描述不一致”类型的负样本。
说明：
- 修改文本中涉及方向的描述（如“南侧”“北侧”“东侧”“西侧”），
  使其与图片<image>中的空间分布相反或不符。
原始文本：
---
{text_chunk}
---
输出：仅返回修改后的文本。"""
    }
    return prompt_templates.get(prompt_type, "")


# ========================== 主程序 ==========================
if __name__ == '__main__':
    doc_path = INPUT_DOCUMENT_PATH
    file_name = os.path.basename(doc_path)
    print(f"\n{'='*80}\n正在处理文件: {file_name}\n{'='*80}")

    text = read_document(doc_path)
    if text.startswith("错误:"):
        print(text)
        exit()

    print(f"文件读取成功，总字数: {len(text)}")

    # 直接使用 read_document 提取的文本作为唯一文本块
    text_chunks = [text]
    print(f"检测到 {len(text_chunks)} 个目标段落。\n")

    all_generated_data = []

    for idx, chunk in enumerate(text_chunks):
        print(f"--- 处理第 {idx+1}/{len(text_chunks)} 段 ---")
        source_label = f"{file_name} 文本块{idx+1}"

        # 遍历4种负样本类型
        for prompt_type in range(1, 5):
            neg_type = ["新增建筑", "漏写建筑", "名称不一致", "方位不一致"][prompt_type - 1]
            
            # 每种类型生成多个样本
            for sample_idx in range(1, SAMPLES_PER_TYPE + 1):
                print(f"  → 生成第{prompt_type}/4类负样本（{neg_type}）- 第{sample_idx}/{SAMPLES_PER_TYPE}个")

                prompt = get_negative_prompt(chunk, source_label, prompt_type)
                msg = [{"role": "user", "content": prompt}]
                result = get_answer(msg)

                if result.strip():
                    all_generated_data.append({
                        "messages": [
                            {
                                "content": f"<image>\n{result.strip()}",
                                "role": "user"
                            },
                            {
                                "content": f"检测到图文不一致问题：{neg_type}。正确的描述应该是：\n{chunk}",
                                "role": "assistant"
                            }
                        ],
                        "images": []
                    })
                else:
                    print(f"  → 负样本生成失败，跳过")
                time.sleep(1)

    # Step 3: 写出结果
    if all_generated_data:
        with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
            json.dump(all_generated_data, f, ensure_ascii=False, indent=4)
        print(f"\n{'='*80}")
        print(f"✅ 处理完成！数据保存至：{OUTPUT_JSON_PATH}")
        print(f"共生成 {len(all_generated_data)} 个对话样本（每个包含负样本user + 正样本assistant）")
        print(f"配置：4种类型 × {SAMPLES_PER_TYPE}个/类型 = {4 * SAMPLES_PER_TYPE} 个对话")
        print("="*80 + "\n")
    else:
        print("\n⚠️ 处理结束，无生成数据。请检查文件或API配置。")
